package com.flame.ai.pdf;

import com.flame.ai.pdf.bean.ImageItem;
import com.flame.ai.pdf.bean.TableItem;
import com.flame.ai.pdf.bean.TextItem;
import com.flame.ai.pdf.config.OllamaProperties;
import com.flame.ai.pdf.service.MineruApiService;
import com.flame.ai.pdf.transformer.DocumentTransformer;
import lombok.Setter;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.document.Document;
import org.springframework.ai.document.DocumentReader;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.core.io.DefaultResourceLoader;
import org.springframework.core.io.FileSystemResource;
import org.springframework.core.io.Resource;
import org.springframework.stereotype.Component;

import java.io.File;
import java.util.List;
import java.util.Map;
import java.util.stream.Stream;

/**
 * @author Jflame
 * @version 1.0.0
 * @since 2025/8/25
 */
@Component
public class PdfDocumentReader implements DocumentReader {
    public static final String METADATA_PAGE_NUMBER = "page_numbers";

    public static final String METADATA_FILE_NAME = "file_name";

    public static final String METADATA_DATA_TYPE = "data_type";

    private final Logger logger = LoggerFactory.getLogger(getClass());

    @Setter
    private Resource file;

    @jakarta.annotation.Resource
    private MineruApiService mineruApiService;

    @jakarta.annotation.Resource
    private OllamaProperties ollamaProperties;


    public void setFile(String filePath) {
        this.file = new DefaultResourceLoader().getResource(filePath);
    }

    public List<Document> get(String filePath) {
        this.file = new DefaultResourceLoader().getResource(filePath);
        this.file = new FileSystemResource(new File(filePath));
        return get();
    }

    public List<Document> get(Resource file) {
        this.file = file;
        return get();
    }

    @Override
    public List<Document> get() {
        if (file == null) {
            logger.warn("No file provided to read.");
            return List.of();
        }
        Map<String, Object> original = mineruApiService.call(file);

        DocumentTransformer transformer = new DocumentTransformer(ollamaProperties.getMultiModel());
        List<Document> textDocuments = transformer.transformTexts((List<TextItem>) original.get("texts"));
        List<Document> imageDocuments = transformer.transformImages((List<ImageItem>) original.get("images"));
        List<Document> tableDocuments = transformer.transformTables((List<TableItem>) original.get("tables"));

        List<Document> documents = Stream.of(textDocuments, imageDocuments, tableDocuments).flatMap(List::stream).toList();
        documents.forEach(document -> document.getMetadata().put(METADATA_FILE_NAME, file.getFilename()));

        logger.info("Successfully process file {}", file.getFilename());
        logger.info("Total documents: {}", documents.size());

        file = null;

        return documents;
    }
}
